package cn.sskxyz.mongodb.multi;

import java.util.Map;

import org.springframework.dao.DataAccessException;
import org.springframework.data.mongodb.core.MongoDatabaseFactorySupport;
import org.springframework.data.mongodb.core.MongoExceptionTranslator;

import com.mongodb.ClientSessionOptions;
import com.mongodb.client.ClientSession;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoDatabase;

import cn.sskxyz.mongodb.multi.dynamic.DynamicDatabaseProvider;

public class MultiMongoDatabaseFactory extends MongoDatabaseFactorySupport<MongoClient> {

	private Map<String, MongoDatabaseFactory> databaseFactory;

	private String primary;

	public MultiMongoDatabaseFactory(MongoClient mongoClient, String databaseName, boolean mongoInstanceCreated,
			Map<String, MongoDatabaseFactory> databaseFactory, String primary) {
		super(mongoClient, databaseName, mongoInstanceCreated, new MongoExceptionTranslator());
		this.databaseFactory = databaseFactory;
		this.primary = primary;
	}

	public Map<String, MongoDatabaseFactory> getDatabaseFactory() {
		return databaseFactory;
	}

	@Override
	protected MongoClient getMongoClient() {
		MongoLinkContext context = MongoDbLinkHolder.get();
		if (context == null || context.getLink() == null) {
			return databaseFactory.get(primary).getMongoClient();
		}

		return databaseFactory.get(context.getLink()).getMongoClient();
	}

	@Override
	public ClientSession getSession(ClientSessionOptions options) {
		return getMongoClient().startSession(options);
	}

	@Override
	public MongoDatabase getMongoDatabase() throws DataAccessException {
		MongoLinkContext context = MongoDbLinkHolder.get();
		if(context == null) {
			return databaseFactory.get(primary).getMongoDatabase();
		}
		
		if (context.getLink() == null ) {
			if (context.getDatabase() == null) {
				return databaseFactory.get(primary).getMongoDatabase();
			}
			return databaseFactory.get(primary).getMongoDatabase(context.getDatabase());
		}

		if (context.getDatabase() == null || context.getDatabase().trim().isEmpty()) {
			return databaseFactory.get(context.getLink()).getMongoDatabase();
		}

		return databaseFactory.get(context.getLink()).getMongoDatabase(context.getDatabase());
	}

	@Override
	protected MongoDatabase doGetMongoDatabase(String dbName) {
		return getMongoClient().getDatabase(dbName);
	}

	@Override
	protected void closeClient() {
		getMongoClient().close();
	}

	public class MongoDatabaseFactory extends MongoDatabaseFactorySupport<MongoClient> {

		public MongoDatabaseFactory(MongoClient mongoClient, String databaseName, boolean mongoInstanceCreated) {
			super(mongoClient, databaseName, mongoInstanceCreated, new MongoExceptionTranslator());
		}

		@Override
		protected MongoClient getMongoClient() {
			return super.getMongoClient();
		}

		@Override
		public ClientSession getSession(ClientSessionOptions options) {
			return getMongoClient().startSession(options);
		}

		@Override
		protected MongoDatabase doGetMongoDatabase(String dbName) {
			return getMongoClient().getDatabase(dbName);
		}

		@Override
		protected void closeClient() {
			getMongoClient().close();
		}

	}

	public class DynamicMongoDatabaseFactory extends MongoDatabaseFactory {

		private DynamicDatabaseProvider dynamicDatabaseProvider;

		public DynamicMongoDatabaseFactory(MongoClient mongoClient, String databaseName, boolean mongoInstanceCreated,
				DynamicDatabaseProvider dynamicDatabaseProvider) {
			super(mongoClient, databaseName, mongoInstanceCreated);
			this.dynamicDatabaseProvider = dynamicDatabaseProvider;
		}

		public String getDefaultDatabaseName() {
			return dynamicDatabaseProvider.getDatabase();
		}
	}
}
